import pandas
import os
from statsmodels.tsa.stattools import adfuller
from statsmodels.tsa.arima.model import ARIMA
import numpy as np
import matplotlib.pyplot as plt

# 读取 CSV 数据
baseDir = './2025_Problem_C_Data'
medals = pandas.read_csv(os.path.join(baseDir, 'summerOly_medal_counts.csv'), encoding='utf-8', low_memory=False)

## 数据预处理
medals['NOC'] = medals['NOC'].str.strip()  # medals 去除 NOC 字段前后的空格

# 以中国为例预测未来的奖牌数量
medals_china = medals[medals['NOC'] == 'China']
print(medals_china)
# 绘制折线图，固定 x 轴，绘制 y 轴为各类奖牌数量的折线图
medals_china.plot(x='Year', y=['Gold', 'Silver', 'Bronze'])
# plt.show()


# 使用 ARIMA 模型对总奖牌数进行预测
data = medals_china['Total'].to_numpy()
result = adfuller(data)
# 平稳性检验
print("ADF Statistic:", result[0])
print("p-value:", result[1])
if result[1] < 0.05:
    print("数据是平稳的")
else:
    print("数据是非平稳的，需要进行差分")
    data_diff = np.diff(data)
    result = adfuller(data_diff)
    print("差分后 ADF Statistic:", result[0])
    print("差分后 p-value:", result[1])

# 定义 ARIMA 模型 (p=1, d=1, q=1)
model = ARIMA(data_diff, order=(1, 1, 1))
model_fit = model.fit()
# 打印模型摘要
print(model_fit.summary())
# 预测未来 1 个时间步
forecast = model_fit.forecast(steps=1)
# 将差分预测值还原到原始数据
forecast = forecast + data[-1]
print("预测结果:", forecast)
